# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of Cybertron package.
#
# The Cybertron is open-source software based on the AI-framework:
# PyTorch (https://pytorch.org/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Metric functions
"""

from typing import Tuple, Optional, Union, List
import numpy as np
from numpy import ndarray

import torch
import torch.nn as nn

from ..utils import get_tensor, get_arguments, GLOBAL_DEVICE


__all__ = [
    'Metric',
    'MaxError',
    'Error',
    'MAE',
    'MSE',
    'MNE',
    'RMSE',
    'Loss',
]

class Metric:
    """Basic class for metrics"""
    def __init__(self, 
                 name,
                 mode="min",
                 index=0,
                 per_atom=False,
                 reduction='mean',
                 aggregate='mean'):
        """
        Args:
            name: 
            mode: 'min' 或 'max', 表示是越小越好还是越大越好
            index: 用于多输出情况，指定计算哪个输出的metric
            per_atom: 是否按原子数平均
            reduction: 输出tensor从(B,Y)到(B,1)的规约方式，可选'mean','sum','none'
            aggregate: 输出tensor从(B,A,...)到(B,A)的聚合方式，可选'mean','sum','none'
        """
        self.name = name
        self.mode = mode
        self.index = index
        self.per_atom = per_atom
        
        # 检查reduction参数
        if reduction is not None:
            if reduction not in ('mean', 'sum', 'none'):
                raise ValueError(f"reduction必须是'mean','sum','none'之一，得到{reduction}")
            self.reduction = None if reduction == 'none' else reduction
        else:
            self.reduction = None
            
        # 检查aggregate参数
        if aggregate is not None:
            if aggregate not in ('mean', 'sum', 'none'):
                raise ValueError(f"aggregate必须是'mean','sum','none'之一，得到{aggregate}")
            self.aggregate = None if aggregate == 'none' else aggregate
        else:
            self.aggregate = None
            
        self.clear()
        self.best_value = float('inf') if self.mode == "min" else float('-inf')
        
    def clear(self):
        """清除状态"""
        self.error_sum = 0
        self.samples_num = 0
        
    def is_better(self, new_value):
        """判断新值是否更好"""
        if new_value is None:
            return False
        if self.mode == "min":
            return new_value < self.best_value
        return new_value > self.best_value
        
    def update_best(self):
        """更新最佳值"""
        new_value = self.eval()
        if self.is_better(new_value):
            self.best_value = new_value
            return True
        return False
    
    def eval(self):
        """返回当前的评估值"""
        return self.error_sum / self.samples_num if self.samples_num > 0 else None
    
    def _select_output(self, outputs, labels):
        """选择要计算的输出和标签"""
        if isinstance(outputs, tuple):
            outputs = outputs[self.index]
        if isinstance(labels, tuple):
            labels = labels[self.index]
        return outputs, labels
    
    def _calc_error(self, predict, label):
        """计算预测值和真实值之间的误差"""
        raise NotImplementedError
        
    def __call__(self, outputs, labels, atom_mask=None):
        """计算指标值
        
        Args:
            outputs: 模型输出，可以是tensor或tuple of tensors
            labels: 真实标签，可以是tensor或tuple of tensors
            atom_mask: 原子mask，形状为(B,A)
            
        Returns:
            float: 计算得到的指标值
        """
        # 选择要计算的输出和标签
        predict, label = self._select_output(outputs, labels)
        
        # 确保数据在正确的设备上
        if atom_mask is not None:
            atom_mask = atom_mask.to(predict.device)
            
        # 计算误差
        error = self._calc_error(predict, label)
        batch_size = error.shape[0]
        
        # 如果维度大于2且需要聚合
        if error.dim() > 2 and self.aggregate is not None:
            # (B,A) <- (B,A,...)
            dims = tuple(range(2, error.dim()))
            if self.aggregate == 'mean':
                error = torch.mean(error, dim=dims)
            else:
                error = torch.sum(error, dim=dims)
                
        # 处理atom_mask
        num_atoms = 0
        total_num = batch_size
        if atom_mask is not None:
            # 计算每个分子的原子数
            num_atoms = torch.count_nonzero(atom_mask, dim=-1) .unsqueeze(-1)
            # total_num = torch.sum(num_atoms)
            # if num_atoms.shape[0] == 1:
            #     total_num *= batch_size
                
            # 应用mask
            if error.shape[1] == atom_mask.shape[1]:
                if error.dim() != atom_mask.dim():
                    # (B,A,...) <- (B,A)
                    new_shape = atom_mask.shape + (1,) * (error.dim() - atom_mask.dim())
                    atom_mask = atom_mask.view(new_shape)
                error = error * atom_mask
                
        # 规约操作
        if self.reduction is not None:
            # (B,1) <- (B,A)
            if self.reduction == 'mean':
                error = torch.mean(error, dim=-1, keepdim=True)
            else:
                error = torch.sum(error, dim=-1, keepdim=True)
                
            # 如果需要按原子数平均
            if self.per_atom and atom_mask is not None:
                error = error / (num_atoms + 1e-8)  # 添加数值稳定性
                
        # 更新状态
        self.error_sum += torch.sum(error).item()
        self.samples_num += total_num
        
        # 返回当前批次的平均误差
        return self.error_sum / self.samples_num if self.samples_num > 0 else 0

class MSE(Metric):
    """均方误差指标"""
    def __init__(self, index=0, per_atom=False, reduction='mean', aggregate='mean'):
        super().__init__("mse", mode="min", index=index, 
                        per_atom=per_atom, reduction=reduction, 
                        aggregate=aggregate)
        
    def _calc_error(self, predict, label):
        return (predict - label) ** 2

class RMSE(Metric):
    """均方根误差指标"""
    def __init__(self, index=0, per_atom=False, reduction='mean', aggregate='mean'):
        super().__init__("rmse", mode="min", index=index,
                        per_atom=per_atom, reduction=reduction,
                        aggregate=aggregate)
        
    def _calc_error(self, predict, label):
        return torch.sqrt((predict - label) ** 2 + 1e-12)  # 添加小量防止零梯度

class MAE(Metric):
    """平均绝对误差指标"""
    def __init__(self, index=0, per_atom=False, reduction='mean', aggregate='mean'):
        super().__init__("mae", mode="min", index=index,
                        per_atom=per_atom, reduction=reduction,
                        aggregate=aggregate)
        
    def _calc_error(self, predict, label):
        return torch.abs(predict - label)

class MaxError(Metric):
    """最大误差指标"""
    def __init__(self, index=0):
        super().__init__("max_error", mode="min", index=index,
                        per_atom=False, reduction=None, aggregate=None)
        self.max_error = 0
        
    def __call__(self, outputs, labels, atom_mask=None):
        # 选择要计算的输出和标签
        predict, label = self._select_output(outputs, labels)
        
        # 计算最大误差范围
        diff = label.reshape(predict.shape) - predict
        max_error = torch.max(diff) - torch.min(diff)
        self.max_error = max(self.max_error, max_error.item())
        
        return self.max_error
    
    def clear(self):
        """清除状态"""
        self.max_error = 0
        self.best_value = float('inf')

class Loss(Metric):
    """损失值指标"""
    def __init__(self):
        super().__init__("loss", mode="min", index=None,
                        per_atom=False, reduction=None, aggregate=None)
        self.clear()
        
    def __call__(self, outputs, labels, atom_mask=None):
        # 对于Loss metric，outputs就是loss值
        loss = outputs
        if isinstance(loss, torch.Tensor):
            loss = loss.item()
            
        # 更新状态
        self.loss_sum += loss
        self.samples_num += 1
        
        # 返回平均损失
        return self.loss_sum / self.samples_num if self.samples_num > 0 else 0
    
    def clear(self):
        """清除状态"""
        self.loss_sum = 0
        self.samples_num = 0
        self.best_value = float('inf')